# encoding: UTF-8

import copy
from collections import Collection

import numpy as np
import pandas as pd

# 分形类型
FRACTAL_NONE = 0
FRACTAL_UP = 1
FRACTAL_DOWN = -1
FRACTAL_UP_CONTINUE = 2
FRACTAL_DOWN_CONTINUE = -2


class FractalData(object):
    '''分形数据'''

    def __init__(self):
        self.datetime = None
        self.ulf = FRACTAL_NONE  # 上下分形标识
        self.uprice = 0
        self.lprice = 0
        self.uprice1 = 0
        self.lprice1 = 0
        self.uperiod = 0
        self.lperiod = 0
        self.uperiod1 = 0
        self.lperiod1 = 0


class FractalCycle(object):
    '''分形循环数据
    4个相邻的分形形成一个分形循环：FC::FCL::(U0,L0,U1,L1)或者FC::FCU::(L0,U0,L1,U1)
    振幅::MAX(U0,U1)-MIN(L0,L1)
    前幅::U0-L0
    后幅::U1-L1
    中幅::FCL(U1-L0),FCU(U0-L1)

    '''

    def __init__(self):
        self.datetime = None


# %load 'http://localhost:8888/notebooks/indicator/fractal.ipynb'
class Fractal(object):
    """
    分形分为上分形和下分形，也可称为顶分形和底分形，记为UF和LF。
    给定周期下的 2*n+1 根相邻的K线，
    $$K=\{k_i, i=0,..., N-1\} = \{(d_i,o_i,h_i,l_i,c_i,v_i),i=1,...,2*n+1\}$$,
    通常n>=2,给出如下定义：
    如果第 n+1 根K线的最高价是这 2*n+1 根K线的最高价,即
    $h_{n+1}=HHV(H,2*n+1)$，称这 2*n+1 根K线形成上分形,记为UF，
    $h_{n+1}$ 称为上分形的顶，记为UPrice，第 2*n+1 根K线称为上分形的确认K线，记为UFOK;
    如果第 n+1 根K线的最低价是这 2*n+1 根K线的最低价,即
    $l_{n+1}=LLV(L,2*n+1)$，称这 2*n+1 根K线形成下分形,记为LF，
    $l_{n+1}$ 称为下分形的底，记为LPrice，第 2*n+1 根K线称为下分形的确认K线，记为LFOK;;
    n称为分形的阶数，满足以下条件
    1、
    2、上、下分形必须轮流出现；
    3、
    """

    def __init__(self, df=None, periods=2, datetime_sgn='datetime', high_sgn='high', low_sgn='low'):
        self.ul = None
        self.uprice = None
        self.datetime_sgn = datetime_sgn
        self.high_sgn = high_sgn
        self.low_sgn = low_sgn
        self.periods = periods
        self.bars = df
        self._len = 0
        if df is not None:
            self._len = len(df.index)
            self.on_bars(self.bars, self.periods, self.datetime_sgn, self.high_sgn, self.low_sgn)

    def on_bars(self, bars, periods=2, datetime_sgn='datetime', high_sgn='high', low_sgn='low'):
        """
        根据K线序列 bars, 计算分形数据
        :param bars: pd.DataFrame
        :param periods: 分形阶数,int
        :param datetime_sgn:
        :param high_sgn:
        :param low_sgn:
        :return:
        """
        dt = list(bars.index)
        high = list(bars[high_sgn])
        low = list(bars[low_sgn])
        self._len = len(high)
        fp = 2 * periods + 1
        bars['ul'] = [0] * self._len  # 上下分形标识，UF＝1，LF＝－1
        bars['fcycle'] = [0] * self._len  # 上下分形循环标识，升LF＝－1
        bars['fbroken'] = [0] * self._len  # 上下分形突破标识，突破上分形＝1，突破下分形＝－1
        bars['uok'] = [0] * self._len  # 上分形确认标识，1＝确认K线
        bars['lok'] = [0] * self._len
        bars['buy'] = [0] * self._len  # 开多标识：1=开多
        bars['sell'] = [0] * self._len  # 平多标识：1=平多
        bars['short'] = [0] * self._len  # 开空标识：1=开空
        bars['cover'] = [0] * self._len  # 平空标识：1=平空
        bars['uperiods'] = [0] * self._len
        bars['lperiods'] = [0] * self._len
        bars['uprice'] = high
        bars['uokprice'] = high
        bars['lprice'] = low
        bars['lokprice'] = low
        bars['uperiods1'] = [0] * self._len
        bars['lperiods1'] = [0] * self._len
        bars['uprice1'] = high
        bars['lprice1'] = low

        i = 0
        mx = high[i]
        mn = low[i]
        ul = 0
        for i in range(1, self._len):
            bar0 = copy.deepcopy(bars.iloc[i - 1])
            bar = copy.deepcopy(bars.iloc[i])
            bar.uprice = bar0.uprice
            bar.uprice1 = bar0.uprice1
            bar.uperiods = bar0.uperiods
            bar.uperiods1 = bar0.uperiods1
            bar.lprice = bar0.lprice
            bar.lprice1 = bar0.lprice1
            bar.lperiods = bar0.lperiods
            bar.lperiods1 = bar0.lperiods1
            bar.fcycle = bar0.fcycle

            if ul == 1:  # 当前处在上分形段
                bar.uperiods = bar0.uperiods + 1
                bar.lperiods = bar0.lperiods + 1
                if mn > low[i]:
                    mn = low[i]
                    mx = high[i]
                    bar.lprice1 = mn
                    bar.lperiods1 = 0
                    bar.uperiods1 = 0
                    if mn < bar.lprice and low[i - 1] > bar.lprice:
                        bar.fbroken = -1
                elif high[i] > mx:  # new high
                    bar.lperiods1 = bar0.lperiods1 + 1
                    mx = high[i]
                    if mx > bar.uprice and high[i - 1] < bar.uprice:
                        bar.fbroken = 1
                    if bar.lperiods1 == periods:  # 下分形确认
                        bar.lok = 1
                        bar.lperiods = 0
                        bar.lokprice = low[i]
                        bar.lprice = mn
                        bar.uperiods1 = 0
                        bar.uprice1 = mx
                        ul = -1
            elif ul == -1:  # 当前处于下分形段
                bar.uperiods = bar0.uperiods + 1
                bar.lperiods = bar0.lperiods + 1
                if high[i] > mx:
                    mn = low[i]
                    mx = high[i]
                    bar.uprice1 = mx
                    bar.lperiods1 = 0
                    bar.uperiods1 = 0
                    if mx > bar.uprice and high[i - 1] < bar.uprice:
                        bar.fbroken = 1

                elif mn > low[i]:  # new low
                    bar.uperiods1 = bar0.uperiods1 + 1
                    mn = low[i]
                    if mn < bar.lprice and low[i - 1] > bar.lprice:
                        bar.fbroken = -1
                    if bar.uperiods1 == periods:  # 下分形确认
                        bar.uok = 1
                        bar.uperiods = 0
                        bar.uokprice = high[i]
                        bar.uprice = mx
                        bar.lperiods1 = 0
                        bar.lprice1 = mn
                        ul = 1
            else:  #
                if mn > low[i]:  # new low
                    bar.uperiods1 = bar0.uperiods1 + 1
                    mn = low[i]
                    if bar.uperiods1 == periods:  # 下分形确认
                        bar.uok = 1
                        bar.uperiods = 0
                        bar.uokprice = high[i]
                        bar.uprice = mx
                        bar.lprice1 = mn
                        bar.lperiods1 = 0
                        ul = 1
                if high[i] > mx:  # new high
                    bar.lperiods1 = bar0.lperiods1 + 1
                    mx = high[i]
                    if bar.lperiods1 == periods:  # 下分形确认
                        bar.lok = 1
                        bar.lperiods = 0
                        bar.lokprice = low[i]
                        bar.uprice1 = mx
                        bar.lprice = mn
                        bar.uperiods1 = 0
                        ul = -1
            bar.buy = 1 if bar.fbroken == 1 else 0
            bar.sell = 1 if bar.fbroken == -1 else 0
            bar.cover = 1 if bar.fbroken == 1 else 0
            bar.short = 1 if bar.fbroken == -1 else 0
            bars.loc[dt[i]] = bar

        self.bars = bars

    def on_bar(self, bars, periods=2, datetime_sgn='datetime', high_sgn='high', low_sgn='low', position=-1):
        hp = bars[high_sgn][position]
        lp = bars[low_sgn][position]


def fractal(dt, high, low, period=2):
    """
    分形
    :param dt: 日期时间（YYYY-MM-DD HH:MM:SS)
    :param high: 最高价
    :param low: 最低价
    :param period: 分形阶数
    :return: 分形序列，FractalData
    """
    fp = 2 * period + 1
    if fp > len(high):
        return None
    ret = []  # OrderedDict()
    fd0 = FractalData()
    i = 0
    fd0.datetime = dt[i]
    fd0.lprice = low[i]
    fd0.uprice = high[i]
    fd0.lprice1 = low[i]
    fd0.uprice1 = high[i]
    fd0.lperiod = 0
    fd0.uperiod = 0
    fd0.lperiod1 = 0
    fd0.uperiod1 = 0
    # ret[fd0.datetime]=fd0
    ret.append(fd0)
    mx = fd0.uprice
    mn = fd0.lprice

    for i in range(1, len(dt)):
        fd = copy.deepcopy(ret[i - 1])
        fd.datetime = dt[i]
        fd.uperiod1 += 1
        fd.lperiod1 += 1
        if high[i] >= mx:  # new high
            mx = high[i]
            if low[i] <= mn:  # new low
                mn = low[i]
                fd.lprice = low[i]
                fd.uprice = high[i]
                fd.lperiod = 0
                fd.uperiod = 0

                ret.append(fd)
                continue
            else:  # only new high
                fd.uprice = mx
                fd.uperiod = 0
                fd.lperiod += 1

                ret.append(fd)
                if fd.lperiod == period:  # find new down-fractal
                    mx = max(high[i - period - 1:i + 1])
                    mn = min(low[i - period - 1:i + 1])
                    for j in range(i - period - 1, i + 1):
                        ret[j].ulf = FRACTAL_DOWN
                        ret[j].lprice = mn
                        ret[j].lprice1 = mn
                        ret[j].lperiod = j - i + period
                        ret[j].lperiod1 = j - i + period
                continue
        elif low[i] <= mn:  # only new low
            mn = low[i]
            fd.lprice = mn
            fd.lperiod = 0
            fd.uperiod += 1

            ret.append(fd)
            if fd.uperiod == period:  # find new up-fractal
                mx = max(high[i - period:i + 1])
                mn = min(low[i - period:i + 1])
                for j in range(i - period, i + 1):
                    ret[j].ulf = FRACTAL_UP
                    ret[j].uprice = mx
                    ret[j].uprice1 = mx
                    ret[j].uperiod = j - i + period
                    ret[j].uperiod1 = j - i + period
            continue
        else:  # no new high and no new low
            fd.lperiod += 1
            fd.uperiod += 1
            ret.append(fd)
            if fd.uperiod == period:  # find new up-fractal
                mx = max(high[i - period:i + 1])
                mn = min(low[i - period:i + 1])
                for j in range(i - period, i + 1):
                    ret[j].ulf = FRACTAL_UP
                    ret[j].uprice = mx
                    ret[j].uprice1 = mx
                    ret[j].uperiod = j - i + period
                    ret[j].uperiod1 = j - i + period

            if fd.lperiod == period:  # find new down-fractal
                row0, row1 = i - period, i + 1
                mn = min(low[row0:row1])
                mx = max(high[row0:row1])
                for j in range(row0, row1):
                    ret[j].ulf = FRACTAL_DOWN
                    ret[j].lprice = mn
                    ret[j].lprice1 = mn
                    ret[j].lperiod = j - i + period
                    ret[j].lperiod1 = j - i + period
            continue
    return ret


def Fractal_T(fr, dt, high, low, period=2):
    """
    分形:根据最新的数据（dt,high,low)和已有的分形数据（fr），计算分形数据
    :param fr:dt之前的分形
    :param dt: 日期时间（YYYY-MM-DD HH:MM:SS)
    :param high: 最新最高价
    :param low: 最新最低价
    :param period: 分形周期
    :return: 分形序列，FractalData
    """

    ret = fr
    i = len(fr)
    fd = copy.deepcopy(ret[i - 1])
    fd.datetime = dt
    mn = fd.lprice
    mx = fd.uprice
    fd.lperiod1 += 1
    fd.uperiod1 += 1
    if high >= mx:
        mx = high
        if low <= mn:
            mn = low
            fd.lprice = mn
            fd.uprice = mx
            fd.lperiod = 0
            fd.uperiod = 0

            ret.append(fd)

        else:
            fd.uprice = mx
            fd.uperiod = 0
            fd.lperiod += 1

            ret.append(fd)
            if fd.lperiod == period:
                for j in range(i - period - 1, i + 1):
                    ret[j].ulf = FRACTAL_DOWN
                    ret[j].lprice = mn
                    ret[j].lprice1 = mn
                    ret[j].lperiod = j - i + period
                    ret[j].lperiod1 = j - i + period

    elif low <= mn:
        mn = low
        fd.lprice = mn
        fd.lperiod = 0
        fd.uperiod += 1
        # ret[fd.datetime] = fd
        ret.append(fd)
        if fd.uperiod == period:
            for j in range(i - period - 1, i + 1):
                ret[j].ulf = FRACTAL_UP
                ret[j].uprice = mx
                ret[j].uprice1 = mx
                ret[j].uperiod = j - i + period
                ret[j].uperiod1 = j - i + period
    else:
        fd.lperiod += 1
        fd.uperiod += 1

        ret.append(fd)
        if fd.uperiod == period:
            for j in range(i - period - 1, i + 1):
                ret[j].ulf = FRACTAL_UP
                ret[j].uprice = mx
                ret[j].uprice1 = mx
                ret[j].uperiod = j - i + period
                ret[j].uperiod1 = j - i + period

        if fd.lperiod == period:
            for j in range(i - period - 1, i + 1):
                ret[j].ulf = FRACTAL_DOWN
                ret[j].lprice = mn
                ret[j].lprice1 = mn
                ret[j].lperiod = j - i + period
                ret[j].lperiod1 = j - i + period

    return ret
